import torch
import numpy as np
from torch import nn 
from torch.autograd import Variable
import torch.nn.functional as F
from datetime import datetime
from torch.utils.data import DataLoader
import Weight_predict

data_train1 = np.loadtxt('data\data1.txt')
data_train2 = np.loadtxt('data\data2.txt')
data_train3 = np.loadtxt('data\data3.txt')

data_train = np.append(data_train1[:,0:-1],data_train2[:,0:-1],axis=0)
data_train = np.append(data_train,data_train3[:,0:-1],axis=0)
data_tr = DataLoader(data_train, 64, shuffle = True)
weight_net = Weight_predict.weight_predict(1)
optim = torch.optim.SGD(weight_net.parameters(), lr=0.001)
cre = nn.MSELoss()
Weight_predict.train(weight_net, data_tr, None, 10, optim, cre)
torch.save(weight_net.state_dict(),'weight_new.pth')

